#include "RHITexture.h"
#include "../../Utils/LogUtil.h"
#include "../../Core/CoreGlobals.h"
#include "../RenderState.h"


RHITexture::~RHITexture()
{
	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		// vkDestroyImageView(GVulkanRHI->GetLogicalDevice(), mImageViews[i], nullptr);
		// vkDestroyImage(GVulkanRHI->GetLogicalDevice(), mImages[i], nullptr);
		// vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mImageMemorys[i], nullptr);
		if (mImageViews[i] != VK_NULL_HANDLE)
			GRenderState->DelayGCImageView(mImageViews[i]);
		if (mImages[i] != VK_NULL_HANDLE)
			GRenderState->DelayGCImage(mImages[i]);
		if (mImageMemorys[i] != VK_NULL_HANDLE)
			GRenderState->DelayGCDeviceMemory(mImageMemorys[i]);
	}
}

VkImageCreateFlags RHITexture::GetImageCreateFlags() const
{
	return mFlags;
}

VkImageType RHITexture::GetImageType() const
{
	return mImageType;
}

VkFormat RHITexture::GetFormat() const
{
	return mFormat;
}

uint32 RHITexture::GetWitdh() const
{
	return mExtent.width;
}

uint32 RHITexture::GetHeight() const
{
	return mExtent.height;
}

uint32 RHITexture::GetDepth() const
{
	return mExtent.height;
}

VkExtent3D RHITexture::GetExtent() const
{
	return mExtent;
}

uint32 RHITexture::GetMipLevels() const
{
	return mMipLevels;
}

uint32 RHITexture::GetArrayLayers() const
{
	return mArrayLayers;
}

VkSampleCountFlagBits RHITexture::GetSamples() const
{
	return mSamples;
}

VkImageTiling RHITexture::GetTiling() const
{
	return mTiling;
}

VkImageUsageFlags RHITexture::GetUsage() const
{
	return mUsage;
}

VkSharingMode RHITexture::GetSharingMode() const
{
	return mSharingMode;
}

std::vector<VkQueueFlagBits> RHITexture::GetQueueFamilys() const
{
	return mQueueFamilys;
}

bool RHITexture::IsInQueueFamily(VkQueueFlagBits queueFlag) const
{
	for (const auto queueFamily : mQueueFamilys)
	{
		if (queueFlag == queueFamily)
			return true;
	}
	return false;
}

VkImageLayout RHITexture::GetImageLayout(uint32 index) const
{
	if (!mIsMultiple)
		return mImageLayouts[0];
	if (index == UINT32_MAX)
		return mImageLayouts[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mImageLayouts[index];
}

VkMemoryPropertyFlags RHITexture::GetMemoryProperty() const
{
	return mMemoryProperty;
}

VkImageViewType RHITexture::GetImageViewType() const
{
	return mViewType;
}

VkComponentMapping RHITexture::GetComponents() const
{
	return mComponents;
}

const VkImageSubresourceRange& RHITexture::GetSubresourceRange() const
{
	return mSubresourceRange;
}

bool RHITexture::IsMultiple() const
{
	return mIsMultiple;
}

bool RHITexture::IsRenderTarget() const
{
	return mIsRenderTarget;
}

VkImage RHITexture::GetImage(uint32 index) const
{
	if (!mIsMultiple)
		return mImages[0];
	if (index == UINT32_MAX)
		return mImages[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mImages[index];
}

VkDeviceMemory RHITexture::GetImageMemory(uint32 index) const
{
	if (!mIsMultiple)
		return mImageMemorys[0];
	if (index == UINT32_MAX)
		return mImageMemorys[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mImageMemorys[index];
}

VkImageView RHITexture::GetImageView(uint32 index) const
{
	if (!mIsMultiple)
		return mImageViews[0];
	if (index == UINT32_MAX)
		return mImageViews[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mImageViews[index];
}

void RHITexture::WriteMemory(void* buffer, uint32 index)
{
	// If you want to copy data, need to be transfer dst bit
	assert(mUsage & VK_IMAGE_USAGE_TRANSFER_DST_BIT);

	// Not support mipmap for now
	assert(mMipLevels == 1);

	VkBuffer stagingBuffer;
	VkDeviceMemory stagingMemory;

	VkDeviceSize dataSize = static_cast<VkDeviceSize>(mExtent.width * mExtent.height * mExtent.depth * GVkFormatSize[mFormat]);

	GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, dataSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

	void* data;
	VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, dataSize, 0, &data));
	memcpy(data, buffer, dataSize);
	vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

	VkImage image = GetImage(index);

	VkImageLayout oldImageLayout = GetImageLayout(index);

	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	GVulkanRHI->TransferLayout(commandBuffer, image, oldImageLayout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);

	GVulkanRHI->CopyBufferToImage(commandBuffer, stagingBuffer, image, mExtent, { mSubresourceRange.aspectMask, 0, 0, mSubresourceRange.layerCount });

	GVulkanRHI->TransferLayout(commandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, oldImageLayout, mSubresourceRange, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT);

	GRenderState->FlushSingleCommand(commandBuffer);

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
}

void RHITexture::ReadMemory(void* buffer, uint32 index)
{
	// If you want to copy data, need to be transfer src bit
	assert(mUsage & VK_IMAGE_USAGE_TRANSFER_SRC_BIT);

	// Not support mipmap for now
	assert(mMipLevels == 1);

	VkBuffer stagingBuffer;
	VkDeviceMemory stagingMemory;

	VkDeviceSize dataSize = static_cast<VkDeviceSize>(mExtent.width * mExtent.height * mExtent.depth * GVkFormatSize[mFormat]);

	GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, dataSize, VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

	VkImage image = GetImage(index);

	VkImageLayout oldImageLayout = GetImageLayout(index);

	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	GVulkanRHI->TransferLayout(commandBuffer, image, oldImageLayout, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);

	GVulkanRHI->CopyImageToBuffer(commandBuffer, image, stagingBuffer, mExtent, { mSubresourceRange.aspectMask , 0, 0, mSubresourceRange.layerCount });

	GVulkanRHI->TransferLayout(commandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, oldImageLayout, mSubresourceRange, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT);

	GRenderState->FlushSingleCommand(commandBuffer);

	void* data;
	VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, dataSize, 0, &data));
	memcpy(buffer, data, dataSize);
	vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

	vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
	vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
}

uint32 RHITexture::GetFrameCount() const
{
	return mIsMultiple ? GRenderState->GetFrameCount() : 1;
}

uint32 RHITexture::GetFrameIndex() const
{
	return mIsMultiple ? GRenderState->GetFrameIndex() : 0;
}

void RHITexture::TransferLayout(VkCommandBuffer commandBuffer, VkImageLayout newLayout, VkImageSubresourceRange subresourceRange, VkPipelineStageFlags srcStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkPipelineStageFlags dstStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkAccessFlags srcAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkAccessFlags dstAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits srcQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits dstQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/)
{
	GVulkanRHI->TransferLayout(commandBuffer, mImages[GetFrameIndex()], mImageLayouts[GetFrameIndex()], newLayout, subresourceRange, srcStageMask, dstStageMask, srcAccessMask, dstAccessMask, srcQueueFlag, dstQueueFlag);
	mImageLayouts[GetFrameIndex()] = newLayout;
}

void RHITexture::TransferLayout(VkImageLayout newLayout, VkImageSubresourceRange subresourceRange, VkQueueFlagBits commandQueue /*= VK_QUEUE_GRAPHICS_BIT*/, VkPipelineStageFlags srcStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkPipelineStageFlags dstStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkAccessFlags srcAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkAccessFlags dstAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits srcQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits dstQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/)
{
	GVulkanRHI->TransferLayout(mImages[GetFrameIndex()], mImageLayouts[GetFrameIndex()], newLayout, subresourceRange, commandQueue, srcStageMask, dstStageMask, srcAccessMask, dstAccessMask, srcQueueFlag, dstQueueFlag);
	mImageLayouts[GetFrameIndex()] = newLayout;
}

void RHITexture::TransferLayout(VkCommandBuffer commandBuffer, VkImageLayout newLayout, VkPipelineStageFlags srcStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkPipelineStageFlags dstStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkAccessFlags srcAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkAccessFlags dstAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits srcQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits dstQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/)
{
	GVulkanRHI->TransferLayout(commandBuffer, mImages[GetFrameIndex()], mImageLayouts[GetFrameIndex()], newLayout, mSubresourceRange, srcStageMask, dstStageMask, srcAccessMask, dstAccessMask, srcQueueFlag, dstQueueFlag);
	mImageLayouts[GetFrameIndex()] = newLayout;
}

void RHITexture::TransferLayout(VkImageLayout newLayout, VkQueueFlagBits commandQueue /*= VK_QUEUE_GRAPHICS_BIT*/, VkPipelineStageFlags srcStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkPipelineStageFlags dstStageMask /*= VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM*/, VkAccessFlags srcAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkAccessFlags dstAccessMask /*= VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits srcQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/, VkQueueFlagBits dstQueueFlag /*= VK_QUEUE_FLAG_BITS_MAX_ENUM*/)
{
	GVulkanRHI->TransferLayout(mImages[GetFrameIndex()], mImageLayouts[GetFrameIndex()], newLayout, mSubresourceRange, commandQueue, srcStageMask, dstStageMask, srcAccessMask, dstAccessMask, srcQueueFlag, dstQueueFlag);
	mImageLayouts[GetFrameIndex()] = newLayout;
}


